import pandas as pd
from base import base
from utils import reduce_mem
class DatasetForLR(base):
    def __init__(self, cfg, need_sample=True, mode="train"):
        super().__init__(cfg)
        self.need_sample = need_sample
        self.mode = mode
        if mode == "train":
            path = self.base_cfg.get('after_train_path')
            data = pd.read_csv(path, index_col=0)

        else:
            path = self.base_cfg.get('after_test_path')
            data = pd.read_csv(path, index_col=0)

        item_related_data = pd.read_csv(self.base_cfg.get('item_related_data'), index_col=0)
        data = data.merge(item_related_data, on='item_id', how='left')
        self.data = reduce_mem(data)
        print(self.data.columns)


    def __call__(self):
        if  self.mode == "test":
            return self.data

        if self.need_sample:
            self.info("train_data_size: {}, pos_data size: {}, neg_data size: {}".format(
                                            self.data.shape[0], 
                                            self.data[self.data['label']==1].shape[0], 
                                            self.data[self.data['label']==0].shape[0]))
            return pd.concat([
                        self.data[self.data['label'] == 0].sample(100000),
                        self.data[self.data['label'] == 1].sample(10000),
                    ])
        else:
            return self.data